﻿using Ndns.Packet.Abstractions;
using Ndns.Packet.Parser;
using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Ndns.Net
{
    
    public class DnsUdpClient : IDisposable
    {
        private IPEndPoint _serverEP;
        private Socket _socket;

        private byte[] buffer = new byte[4096];
        private volatile int _isBusy = 0;
        private IEnumerator<ushort> _sequence;
        
        private bool _useRandomTransactionId;
        public bool UseRandomTransactionId
        {
            get => _useRandomTransactionId;
            set
            {
                if (_useRandomTransactionId != value)
                {
                    _sequence.Dispose();
                    if (value)
                        _sequence = TransactionIdGenerator.GetRandomTransactionId().GetEnumerator();
                    else
                        _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
                    _useRandomTransactionId = value;
                }
            }
        }

        public IPEndPoint ServerEndPoint => _serverEP;

        public bool EnableCompression { get; set; } = true;
        public bool AllowTruncation { get; set; } = true;

        private DnsPacketParser _packetParser;

        public DnsUdpClient()
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = new DnsPacketParser();
        }
        public DnsUdpClient(DnsPacketParser packetParser)
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = packetParser ?? new DnsPacketParser();
        }



        public Task ConnectAsync(IPEndPoint serverEP)
        {
            if (_socket != null)
                _socket.Dispose();
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));

            _serverEP = serverEP ?? throw new ArgumentNullException(nameof(serverEP));
            InternalInitSocket();

            return Task.CompletedTask;
        }
        private void InternalInitSocket()
        {
            _socket = new Socket(_serverEP.AddressFamily, SocketType.Dgram, ProtocolType.Udp);
            _socket.Bind(new IPEndPoint(_serverEP.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, 0));
        }

        
        public async Task<DnsPacket> SendAsync(DnsQuery query, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (query == null)
                throw new ArgumentNullException(nameof(query));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));
            if (_socket == null)
                throw new InvalidOperationException();
            if (_isBusy == 1)
                throw new InvalidOperationException();
            if (1 == Interlocked.Exchange(ref _isBusy, 1))
                throw new InvalidOperationException();
            
            _sequence.MoveNext();
            DnsPacket request = new DnsPacket(_sequence.Current, DnsPacketFlag.RecursionDesired);
            request.Queries.Add(query);

            try
            {
                return await InternalSendAsync(request, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                _isBusy = 0;
            }
        }

        public async Task<DnsPacket> SendAsync(DnsPacket request, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (request == null)
                throw new ArgumentNullException(nameof(request));
            if (disposedValue)
                throw new ObjectDisposedException(nameof(DnsUdpClient));
            if (_socket == null)
                throw new InvalidOperationException();
            if (_isBusy == 1)
                throw new InvalidOperationException();
            if (1 == Interlocked.Exchange(ref _isBusy, 1))
                throw new InvalidOperationException();
            
            try
            {
                return await InternalSendAsync(request, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                _isBusy = 0;
            }
        }

        private async Task<DnsPacket> InternalSendAsync(DnsPacket request, CancellationToken cancellationToken)
        {
            using (cancellationToken.Register(() => _socket.Dispose()))
            {
                try
                {
                    //接收数据
                    var taskReceive = _socket.ReceiveFromAsync(new ArraySegment<byte>(buffer), SocketFlags.None, _serverEP).ConfigureAwait(false);
                    //发送数据包
                    var reqPacketBuff = request.GetBytes(AllowTruncation, EnableCompression);
                    await _socket.SendToAsync(new ArraySegment<byte>(reqPacketBuff), SocketFlags.None, _serverEP).ConfigureAwait(false);
                    //获取接收的数据
                    var received = await taskReceive;
                    using (var ms = new MemoryStream(buffer,0 , received.ReceivedBytes, false))
                    {
                        if (!_packetParser.TryParse(ms, out var packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                        return packet;
                    }
                }
                catch (ObjectDisposedException)
                {
                    InternalInitSocket();
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }

        }
        
        public static async Task<DnsPacket> SendAsync(DnsPacket request, IPEndPoint server, CancellationToken cancellationToken = default(CancellationToken))
        {
            var _parser = new DnsPacketParser();
            //准备Socket
            using (var udpSocket = new Socket(server.AddressFamily, SocketType.Dgram, ProtocolType.Udp))
            using (cancellationToken.Register(() => udpSocket.Dispose()))
            {
                udpSocket.Bind(new IPEndPoint(server.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, 0));
                try
                {
                    //接收数据
                    var buffer = new byte[4096];
                    var taskReceive = udpSocket.ReceiveFromAsync(new ArraySegment<byte>(buffer), SocketFlags.None, server).ConfigureAwait(false);
                    //发送数据包
                    var reqPacketBuff = request.GetBytes(true, true);
                    await udpSocket.SendToAsync(new ArraySegment<byte>(reqPacketBuff), SocketFlags.None, server).ConfigureAwait(false);
                    //获取接收的数据
                    var received = await taskReceive;
                    using (var ms = new MemoryStream(buffer, 0, received.ReceivedBytes, false))
                    {
                        if (!_parser.TryParse(ms, out var packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                        return packet;
                    }
                }
                catch (ObjectDisposedException)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }
        }

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                    _socket?.Dispose();
                    _sequence?.Dispose();
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。
                _socket = null;
                _sequence = null;

                disposedValue = true;
            }
        }


        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
        }
        #endregion
    }
}
